import google.generativeai as genai
import os
import logging
import time # 导入 time 模块以添加延迟

# --- 严重安全警告 ---
# 您提供的 API 密钥 'AIzaSyBHLohL3JURhWWoWa28JVqHeFXDK7hNgrQ' 已公开，极度不安全！
# 请立即前往 Google Cloud Console 或 MakerSuite 停用此密钥。
# 强烈建议将新的、安全的 API 密钥存储在环境变量中，例如：
# API_KEY = os.environ.get("GOOGLE_API_KEY")
# 然后在您的运行环境中设置 GOOGLE_API_KEY 变量。
#
# 为了运行您要求的测试，暂时使用您提供的密钥，但这存在巨大风险。
API_KEY = "AIzaSyBHLohL3JURhWWoWa28JVqHeFXDK7hNgrQ" # <--- 极度不安全！请立即替换和撤销！

if not API_KEY:
    print("错误：环境变量 'GOOGLE_API_KEY' 未设置。")
    print("请在运行脚本前通过 'export GOOGLE_API_KEY=\"YOUR_API_KEY\"' 命令设置它。")
    exit()

try:
    genai.configure(api_key=API_KEY)
    print("✅ Google Generative AI SDK 配置成功。")
except Exception as e:
    print(f"❌ 错误：配置 SDK 失败: {e}")
    exit()

# --- 2. 准备要测试的模型列表 ---
# 这是您提供的完整列表
models_to_test = [
 
    "models/gemini-2.5-flash-preview-05-20", "models/gemini-2.5-flash", "models/gemini-2.5-flash-lite-preview-06-17",
    "models/gemini-2.5-pro-preview-05-06", "models/gemini-2.5-pro",
    "models/gemini-2.0-flash-exp", "models/gemini-2.0-flash", "models/gemini-2.0-flash-001",
]

# 用于测试的简单提示词
test_prompt = "用一句话简单介绍一下你自己。"

print(f"\n准备测试 {len(models_to_test)} 个模型...\n")

# --- 3. 逐一测试并记录结果 ---
available_models = []
unavailable_models = []
errored_models = []

for model_name in models_to_test:
    print(f"--- 正在测试: {model_name} ---")
    try:
        model = genai.GenerativeModel(model_name)
        response = model.generate_content(test_prompt)
        
        # 检查响应是否有效
        if response and response.text:
            print(f"✅ 成功! 模型可用。")
            available_models.append(model_name)
        elif response and response.prompt_feedback:
            # 可能是因为安全过滤等原因，但模型本身是可调用的
            print(f"⚠️ 注意: 模型可用，但因安全过滤无文本返回。 Feedback: {response.prompt_feedback}")
            available_models.append(f"{model_name} (with content filter)")
        else:
            # 其他未知成功情况
            print("❔ 成功调用但响应为空，计入可用列表。")
            available_models.append(f"{model_name} (empty response)")

    except Exception as e:
        error_message = str(e)
        if "is not found" in error_message or "404" in error_message:
            print(f"❌ 不可用 (Not Found): 该模型已下线或不存在。")
            unavailable_models.append(model_name)
        else:
            # 记录其他类型的错误
            print(f"⚠️ 出错: {type(e).__name__} - {error_message[:150]}...")
            errored_models.append(f"{model_name} (Error: {type(e).__name__})")
        
    # 在每个请求之间加入短暂延迟，避免触发频率限制
    time.sleep(2)

# --- 4. 打印最终的汇总报告 ---
print("\n" + "="*50)
print("          模型可用性测试报告")
print("="*50)

print(f"\n✅ 可用模型 ({len(available_models)} 个):")
if available_models:
    for m in available_models:
        print(f"  - {m}")
else:
    print("  - 无")

print(f"\n❌ 不可用 / 已下线模型 ({len(unavailable_models)} 个):")
if unavailable_models:
    for m in unavailable_models:
        print(f"  - {m}")
else:
    print("  - 无")

print(f"\n⚠️ 测试时出错的模型 ({len(errored_models)} 个):")
if errored_models:
    for m in errored_models:
        print(f"  - {m}")
else:
    print("  - 无")

print("\n" + "="*50)
print("\n报告解读:")
print("  - '✅ 可用模型' 列表中的模型，在您的免费额度内都可以放心使用。")
print("  - 对于您的 Flask 应用，强烈建议从 '✅ 可用模型' 列表中选择一个稳定的（不含 preview/exp 后缀的）模型，如 'models/gemini-1.5-flash'。")